import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from matplotlib import gridspec
import matplotlib.animation as animation
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d.proj3d import proj_transform
from mpl_toolkits.mplot3d.axes3d import Axes3D


class Arrow3D(FancyArrowPatch):
    def __init__(self, x, y, z, dx, dy, dz, *args, **kwargs):
        super().__init__((0, 0), (0, 0), *args, **kwargs)
        self._xyz = (x, y, z)
        self._dxdydz = (dx, dy, dz)

    def draw(self, renderer):
        x1, y1, z1 = self._xyz
        dx, dy, dz = self._dxdydz
        x2, y2, z2 = (x1 + dx, y1 + dy, z1 + dz)

        xs, ys, zs = proj_transform((x1, x2), (y1, y2), (z1, z2), renderer.M)
        self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
        super().draw(renderer)


def _arrow3D(ax, x, y, z, dx, dy, dz, *args, **kwargs):
    '''Add an 3d arrow to an `Axes3D` instance.'''

    arrow = Arrow3D(x, y, z, dx, dy, dz, *args, **kwargs)
    ax.add_artist(arrow)


setattr(Axes3D, 'arrow3D', _arrow3D)


class PlotScene:
    def __init__(self):
        self.c_ang = np.linspace(0, 2 * np.pi, 100)
        self.c_x = np.cos(self.c_ang)
        self.c_y = np.sin(self.c_ang)
        self.c_z = np.zeros_like(self.c_x)
        self.alpha1 = 0.15
        self.alpha2 = 0.05
        self.plane_x = np.array([[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]])
        self.plane_y = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])
        self.plane_z = np.zeros((3, 3))

    def ket(self, ax, fontsize=18):
        ax.text(0, 0, 1.1, r'$\hat{{z}}=\left|0\right>$', fontsize=fontsize)
        ax.text(0, 1.2, -0.1, r'$\hat{{y}}$', fontsize=fontsize)
        ax.text(1.0, 0, -0.2, r'$\hat{{x}}$', fontsize=fontsize)

    def circuit_yz(self, ax):
        ax.plot(self.c_z, self.c_x, self.c_y, 'k')

    def circuit_xy(self, ax, style='k'):
        ax.plot(self.c_x, self.c_y, self.c_z, style)

    def rect(self, ax, center_x, center_y, width, height, edgecolor='k'):
        ax.add_patch(
            plt.Rectangle((center_x - width / 2, center_y - height / 2),
                          width,
                          height,
                          facecolor='w',
                          edgecolor=edgecolor,
                          linewidth=2))

    def xy(self, ax):
        ax.plot_wireframe(self.plane_x,
                          self.plane_y,
                          self.plane_z,
                          color='g',
                          alpha=self.alpha1)
        ax.plot_surface(self.plane_x,
                        self.plane_y,
                        self.plane_z,
                        color='g',
                        alpha=self.alpha2)

    def xz(self, ax):
        ax.plot_wireframe(self.plane_x,
                          self.plane_z,
                          self.plane_y,
                          color='r',
                          alpha=self.alpha1)
        ax.plot_surface(self.plane_x,
                        self.plane_z,
                        self.plane_y,
                        color='r',
                        alpha=self.alpha2)

    def yz(self, ax):
        ax.plot_wireframe(self.plane_z,
                          self.plane_x,
                          self.plane_y,
                          color='b',
                          alpha=self.alpha1)
        ax.plot_surface(self.plane_z,
                        self.plane_x,
                        self.plane_y,
                        color='b',
                        alpha=self.alpha2)

    def set_view(self, ax):
        ax.set_xlim3d(-1, 1)
        ax.set_ylim3d(-1, 1)
        ax.set_zlim3d(-1, 1)
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.set_box_aspect([1, 1, 1])
        ax.view_init(elev=0, azim=0)

    def text(self, ax, x, y, text, fs, **kargs):
        return ax.text(x,
                       y,
                       text,
                       fontsize=fs,
                       verticalalignment='center',
                       horizontalalignment='left',
                       **kargs)


def update_lines(num, data_line, angle_data, ansatz_data, line, vec, point,
                 tx1, tx2, line1, line2, line3):
    # line.set_data(data_line[0:2, :num])
    # line.set_3d_properties(data_line[2, :num])
    vec.set_data(np.array([[0, data_line[0, num]], [0, data_line[1, num]]]))
    vec.set_3d_properties(np.array([0, data_line[2, num]]))
    line1.set_data(np.array([[0, 0], [0, 0]]))
    line1.set_3d_properties(np.array([data_line[2, num], 0]))
    line2.set_data(np.array([[data_line[0, num], 0], [0, 0]]))
    line2.set_3d_properties(np.array([0, 0]))
    line3.set_data(np.array([[0, 0], [data_line[1, num], 0]]))
    line3.set_3d_properties(np.array([0, 0]))
    theta, psi = angle_data[num][0], angle_data[num][1]
    s1 = str(np.cos(theta))[:4]
    s2 = str(psi)[:4]
    s3 = str(np.sin(theta))[:4]
    ax2.set_title(
        r'$\left|\varphi\right>_f={}\left|0\right>+e^{{{}i}}{}\left|1\right>$'.
        format(s1, s2, s3),
        fontsize=18,
        y=0.8)
    point._offsets3d = data_line[:, num][:, None]
    tx1.set_text('RX({:.2})'.format(round(ansatz_data[num, 2], 2)))
    tx2.set_text('RY({:.2})'.format(round(ansatz_data[num, 3], 2)))


data = np.loadtxt('data.csv')
params = np.loadtxt('params.csv')
final_angle = np.loadtxt('final_angle.csv')

fig = plt.figure()
# gs = gridspec.GridSpec(2, 1, height_ratios=[7, 1])
# ax1 = plt.subplot(gs[0], projection='3d')
ax1 = fig.add_axes([0, 0.2, 1, 0.8], projection='3d')
ax2 = fig.add_axes([0.15, 0.0, 0.7, 0.3])

ax1.arrow3D(0,
            0,
            -1,
            0,
            0,
            2,
            arrowstyle='-|>',
            mutation_scale=20,
            linestyle='dashed')
ax1.arrow3D(0,
            -1,
            0,
            0,
            2,
            0,
            arrowstyle='-|>',
            mutation_scale=20,
            linestyle='dashed')
ax1.arrow3D(-1,
            0,
            0,
            2,
            0,
            0,
            arrowstyle='-|>',
            mutation_scale=20,
            linestyle='dashed')
scene = PlotScene()
scene.xy(ax1)
scene.yz(ax1)
scene.xz(ax1)
scene.circuit_xy(ax1, 'k--')
scene.circuit_yz(ax1)
scene.set_view(ax1)
scene.ket(ax1)
line = ax1.plot(data[0, 0:1], data[1, 0:1], data[2, 0:1], 'k--',
                linewidth=1)[0]
line1 = ax1.plot([0, data[0, 0]], [0, data[1, 0]], [data[2, 0], data[2, 0]],
                 'g',
                 linewidth=3)[0]
line2 = ax1.plot([data[0, 0], data[0, 0]], [0, data[1, 0]], [0, data[2, 0]],
                 'g',
                 linewidth=3)[0]
line3 = ax1.plot([0, data[0, 0]], [data[1, 0], data[1, 0]], [0, data[2, 0]],
                 'g',
                 linewidth=3)[0]
ax1.plot([0, data[0, 0]], [0, data[1, 0]], [0, data[2, 0]], 'k',
         linewidth=3)[0]
ax1.scatter([data[0, 0]], [data[1, 0]], [data[2, 0]], s=50, marker='o', c='k')
vec = ax1.plot([0, data[0, 0]], [0, data[1, 0]], [0, data[2, 0]],
               'r',
               linewidth=3)[0]
point = ax1.scatter([data[0, 0]], [data[1, 0]], [data[2, 0]],
                    s=50,
                    marker='o',
                    c='r')
ax1.xaxis.set_major_locator(ticker.MultipleLocator(0.5))
ax1.yaxis.set_major_locator(ticker.MultipleLocator(0.5))
ax1.set_axis_off()
# ax2 = plt.subplot(gs[1])

ax2.axis('equal')
ax2.plot([-0.18, 0.19], [0, 0], 'k', linewidth=2, zorder=-1)
bx = -0.2
by = 0.0
h = w = 0.045
scene.text(ax2, bx - 0.01, by, r'$\left|0\right>$', 20)
bx += 0.08
scene.rect(ax2, bx, by, h, w)
scene.text(ax2,
           bx - h / 2 * 0.8,
           by,
           'RX({:.3})'.format(round(params[0, 0], 3)),
           9,
           fontweight='bold')
scene.text(ax2, bx + 0.01, by - 0.04, 'Encoder', 16)
bx += 0.06
scene.rect(ax2, bx, by, h, w)
scene.text(ax2,
           bx - h / 2 * 0.8,
           by,
           'RY({:.3})'.format(round(params[0, 1], 3)),
           9,
           fontweight='bold')
bx += 0.10
scene.rect(ax2, bx, by, h, w)
a1 = scene.text(ax2,
                bx - h / 2 * 0.8,
                by,
                'RX({:.3})'.format(round(params[0, 2], 3)),
                9,
                fontweight='bold')
scene.text(ax2, bx + 0.01, by - 0.04, 'Ansatz', 16)
bx += 0.06
scene.rect(ax2, bx, by, h, w)
a2 = scene.text(ax2,
                bx - h / 2 * 0.8,
                by,
                'RY({:.3})'.format(round(params[0, 3], 3)),
                9,
                fontweight='bold')
bx += 0.08
scene.text(ax2, bx + 0.015, by, r'$\left|\varphi\right>_f$', 20)
ax2.axes.get_yaxis().set_visible(False)
ax2.axes.get_xaxis().set_visible(False)
for key, spine in ax2.spines.items():
    spine.set_visible(False)
line_ani = animation.FuncAnimation(fig,
                                   update_lines,
                                   data.shape[1],
                                   fargs=(data, final_angle, params, line, vec,
                                          point, a1, a2, line1, line2, line3),
                                   interval=10,repeat=False)
plt.show()